{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# {class}`tvm.arith.analyzer.ConstIntBound`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/read/arith\n"
     ]
    }
   ],
   "source": [
    "%cd ..\n",
    "import testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import contextlib\n",
    "import tvm\n",
    "import tvm.testing\n",
    "\n",
    "from tvm import te\n",
    "from tvm.arith import ConstIntBound\n",
    "\n",
    "NEG_INF = ConstIntBound.NEG_INF\n",
    "POS_INF = ConstIntBound.POS_INF\n",
    "\n",
    "class TestCase:\n",
    "    def __init__(self, expr, expected_bounds, known_bounds=None, constraint=None):\n",
    "        self.expr = expr\n",
    "        self.expected_bounds = expected_bounds\n",
    "        if known_bounds is None:\n",
    "            self.known_bounds = {}\n",
    "        else:\n",
    "            self.known_bounds = known_bounds\n",
    "\n",
    "        self.constraint = constraint\n",
    "\n",
    "    @property\n",
    "    def __name__(self):\n",
    "        return str(self.expr)\n",
    "\n",
    "\n",
    "class BaseCompare:\n",
    "    def test_const_bounds(self, test_case):\n",
    "        analyzer = tvm.arith.Analyzer()\n",
    "\n",
    "        for var, bounds in test_case.known_bounds.items():\n",
    "            analyzer.update(var, ConstIntBound(*bounds))\n",
    "\n",
    "        with contextlib.ExitStack() as stack:\n",
    "            if test_case.constraint is not None:\n",
    "                stack.enter_context(analyzer.constraint_scope(test_case.constraint))\n",
    "\n",
    "            bounds = analyzer.const_int_bound(test_case.expr)\n",
    "\n",
    "        if test_case.expected_bounds[0] is None:\n",
    "            assert bounds.max_value == test_case.expected_bounds[1]\n",
    "        elif test_case.expected_bounds[1] is None:\n",
    "            assert bounds.min_value == test_case.expected_bounds[0]\n",
    "        else:\n",
    "            assert (bounds.min_value, bounds.max_value) == test_case.expected_bounds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class TestDataType(BaseCompare):\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(te.var(\"x\", dtype=\"int64\"), (NEG_INF, POS_INF)),\n",
    "        TestCase(te.var(\"x\", dtype=\"int8\"), (-128, 127)),\n",
    "        TestCase(te.var(\"x\", dtype=\"uint8\"), (0, 255)),\n",
    "        TestCase(te.size_var(\"x\", dtype=\"int32\"), (0, POS_INF)),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestCastBound(BaseCompare):\n",
    "    x = te.var(\"x\", dtype=\"int8\")\n",
    "    tmod = tvm.tir.truncmod\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tmod(x, 3).astype(\"uint32\"), (0, 2)),\n",
    "        TestCase(tmod(x, 3).astype(\"float32\").astype(\"int32\"), (-2, 2)),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestAddSubBound(BaseCompare):\n",
    "    x = te.var(\"x\", \"int64\")\n",
    "    y = te.var(\"y\", \"int64\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(x + y, (NEG_INF, POS_INF)),\n",
    "        TestCase(x + y, (1, 14), known_bounds={x: (0, 4), y: (1, 10)}),\n",
    "        TestCase(x - y, (-10, 3), known_bounds={x: (0, 4), y: (1, 10)}),\n",
    "        TestCase(x - y, (-10, POS_INF), known_bounds={x: (0, POS_INF), y: (1, 10)}),\n",
    "        TestCase(1 - x, (NEG_INF, 1), known_bounds={x: (0, POS_INF), y: (1, 10)}),\n",
    "    )\n",
    "\n",
    "\n",
    "@pytest.mark.xfail(reason=\"Not currently supported\")\n",
    "class TestBoundsUsingReciprocals(BaseCompare):\n",
    "    \"\"\"Special handling for differences of reciprocals\n",
    "\n",
    "    These terms can appear when comparing the number of operations for\n",
    "    different orderings of matrix multiplications, with A, B, and C\n",
    "    known to be positive values.\n",
    "\n",
    "    In these cases, comparing `(A+B)*C < A*B` is equivalent to\n",
    "    `1/A + 1/B < 1/C`.  Working in terms of the reciprocals\n",
    "    allows the ConstIntBound analyzer to provide a tighter\n",
    "    bound for these differences than would otherwise be\n",
    "    available.\n",
    "\n",
    "    For `(A+B)*C - A*B`, the normal bottom-up integer bounds are unable to\n",
    "    provide the bounds required to provide these inequalities, because they\n",
    "    treat the terms as uncorrelated.  That is, they assume that `(A+B)*C` may\n",
    "    achieve its minimum while `A*B` simultaneously achieves its maximum.\n",
    "    \"\"\"\n",
    "\n",
    "    A, B, C = [te.var(letter, \"int64\") for letter in \"ABC\"]\n",
    "\n",
    "    symmetric_bounds = {A: (1, 4095), B: (1, 4095), C: (2048, 2048)}\n",
    "    asymmetric_bounds = {A: (1, 1024), B: (1, POS_INF), C: (2048, 2048)}\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase((A + B) * C - A * B, (2048, None), known_bounds=symmetric_bounds),\n",
    "        TestCase((A + B) * C - B * A, (2048, None), known_bounds=symmetric_bounds),\n",
    "        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),\n",
    "        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=symmetric_bounds),\n",
    "        TestCase((A + B) * C - A * B, (2048, None), known_bounds=asymmetric_bounds),\n",
    "        TestCase((A + B) * C - B * A, (2048, None), known_bounds=asymmetric_bounds),\n",
    "        TestCase(A * B - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),\n",
    "        TestCase(B * A - (A + B) * C, (None, -2048), known_bounds=asymmetric_bounds),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestMulBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(x * y + 20, (0, 60), {x: (-2, 4), y: (4, 10)}),\n",
    "        TestCase(x * y, (-32, 24), {x: (-3, 4), y: (-8, 2)}),\n",
    "        TestCase(x * y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-8, 2)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestTruncDivBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    expr = tvm.tir.truncdiv(x, y)\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(expr, (-2, None), {x: (-9, 4), y: (4, 10)}),\n",
    "        TestCase(expr, (-4, 9), {x: (-9, 4), y: (-2, 0)}),\n",
    "        TestCase(expr, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),\n",
    "        TestCase(expr, (-9, 9), {x: (-9, 4), y: (-4, 12)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestTruncModBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    expr = tvm.tir.truncmod(x, y)\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(expr, (-9, 4), {x: (-9, 4), y: (4, 10)}),\n",
    "        TestCase(expr, (-9, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),\n",
    "        TestCase(expr, (0, 9), {x: (1, POS_INF), y: (4, 10)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestFloorDivBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "    ux = te.var(\"x\", dtype=\"uint32\")\n",
    "    uy = te.var(\"y\", dtype=\"uint32\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(x // y, (-9 // 4, None), {x: (-9, 4), y: (4, 10)}),\n",
    "        TestCase(x // y, (-4, 9), {x: (-9, 4), y: (-2, 0)}),\n",
    "        TestCase(x // y, (NEG_INF, POS_INF), {x: (NEG_INF, 4), y: (-2, 1)}),\n",
    "        TestCase(x // y, (-9, 9), {x: (-9, 4), y: (-4, 12)}),\n",
    "        TestCase(ux // uy, (0, 4), {ux: (1, 4), uy: (0, 12)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestFloorModBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(x % y, (0, 9), {x: (-9, 4), y: (4, 10)}),\n",
    "        TestCase(x % y, (0, 9), {x: (NEG_INF, POS_INF), y: (4, 10)}),\n",
    "        TestCase(x % y, (0, 9), {x: (1, POS_INF), y: (4, 10)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestMinMaxBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tvm.te.min(x, y), (-9, 10), {x: (-9, 11), y: (4, 10)}),\n",
    "        TestCase(tvm.te.min(x, y), (NEG_INF, 10), {x: (NEG_INF, POS_INF), y: (4, 10)}),\n",
    "        TestCase(tvm.te.max(x, y), (4, POS_INF), {x: (NEG_INF, POS_INF), y: (4, 10)}),\n",
    "        TestCase(tvm.te.max(x, y), (4, POS_INF), {x: (1, POS_INF), y: (4, 10)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestSelectBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(\n",
    "            tvm.tir.Select(x > 1, (y < 0).astype(\"int32\"), y + 1),\n",
    "            (0, 11),\n",
    "            {x: (-9, 11), y: (4, 10)},\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestShiftAndBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(x >> y, (-3, 2), {x: (-9, 11), y: (2, 10)}),\n",
    "        TestCase(x & y, (0, 10), {x: (-9, 11), y: (2, 10)}),\n",
    "        TestCase(x & y, (0, 10), {x: (10, 11), y: (2, 10)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestMixIndexBound(BaseCompare):\n",
    "    x, y = te.var(\"x\"), te.var(\"y\")\n",
    "    tdiv = tvm.tir.truncdiv\n",
    "    tmod = tvm.tir.truncmod\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tmod(x, 8) + tdiv(x, 8) * 8, (0, 24 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),\n",
    "        TestCase(y + x * 3, (0, 24 * 3 - 1), {x: (0, 24 - 1), y: (0, 3 - 1)}),\n",
    "        TestCase(\n",
    "            tmod(x, 7) + tdiv(x, 7) * 7, (0, (23 // 7) * 7 + 6), {x: (0, 24 - 1), y: (0, 3 - 1)}\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestLetBound(BaseCompare):\n",
    "    x = te.var(\"x\")\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tvm.tir.Let(x, 1, x + 1), (2, 2)),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestFloorModNegativeDivisor(BaseCompare):\n",
    "    flm, fld = tvm.te.floormod, tvm.te.floordiv\n",
    "    a, b = te.var(\"a\"), te.var(\"b\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(a % b, (-4, 6), {a: (0, 6), b: (-5, 7)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestDivModAssumeNoZeroDivisor(BaseCompare):\n",
    "    \"\"\"Divmod non negative expression makes assumption that divide by\n",
    "    zero won't occur this assumption is important to get best result\n",
    "    from symbolic shape programs\n",
    "    \"\"\"\n",
    "\n",
    "    a, b = te.var(\"a\"), te.var(\"b\")\n",
    "\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(a // b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),\n",
    "        TestCase(a % b, (0, 6), {a: (0, 6), b: (0, POS_INF)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestMultipleCondition(BaseCompare):\n",
    "    a = te.var(\"a\")\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(\n",
    "            a % 58 - 1,\n",
    "            (0, None),\n",
    "            known_bounds={a: (0, 128)},\n",
    "            constraint=tvm.tir.all(1 <= a % 58, a % 58 < 57),\n",
    "        ),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestBroadcastBound(BaseCompare):\n",
    "    a = te.var(\"a\")\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tvm.tir.Broadcast(a, 4), (0, 128), {a: (0, 128)}),\n",
    "    )\n",
    "\n",
    "\n",
    "class TestRampBound(BaseCompare):\n",
    "    a = te.var(\"a\")\n",
    "    test_case = tvm.testing.parameter(\n",
    "        TestCase(tvm.tir.Ramp(a, 2, 4) + 2, (2, 128 + 2 * 3 + 2), {a: (0, 128)}),\n",
    "    )\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
